import torch
from torchvision import datasets, models, transforms
import os

def train(model_path, dataloaders, class_names):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(device)

    # 加载预训练模型
    model = models.mobilenet_v2(pretrained=True)
    model.to(device)

    # 冻结预训练模型的所有层
    for param in model.parameters():
        param.requires_grad = False

    # 替换最后一层
    num_ftrs = model.classifier[1].in_features  # 注意这里的索引应该是1
    model.classifier[1] = torch.nn.Linear(num_ftrs, len(class_names)).to(device)  # 确保新层也在目标设备上

    # 定义损失函数和优化器
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9)  # 只优化最后一层

    # 微调模型
    num_epochs = 30
    best_acc = 0.0

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        running_corrects = 0

        for inputs, labels in dataloaders['train']:
            inputs = inputs.to(device)
            labels = labels.to(device)

            optimizer.zero_grad()

            # 前向传播
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)

            # 计算损失
            loss = criterion(outputs, labels)

            # 反向传播和优化
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * inputs.size(0)
            running_corrects += (preds == labels).sum().item()

        # 计算每个epoch的损失和准确率
        epoch_loss = running_loss / len(dataloaders['train'].dataset)
        epoch_acc = running_corrects / len(dataloaders['train'].dataset)

        print('Epoch [{}/{}], Loss: {:.4f}, Acc: {:.4f}'.format(epoch + 1, num_epochs, epoch_loss, epoch_acc))

        if epoch_acc > best_acc:
            best_acc = epoch_acc
            best_model_wts = model.state_dict()

    # 保存最佳模型权重
    if not os.path.exists(model_path):
        os.makedirs(model_path)
    torch.save(best_model_wts, os.path.join(model_path, 'best_model.pth'))

def test(model_path, dataloaders, class_names):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 加载预训练模型
    model = models.mobilenet_v2(pretrained=True)
    model.to(device)

    # 替换最后一层以匹配类别数量
    num_ftrs = model.classifier[1].in_features
    model.classifier[1] = torch.nn.Linear(num_ftrs, len(class_names)).to(device)

    # 加载保存的模型权重
    model.load_state_dict(torch.load(os.path.join(model_path, 'best_model.pth')))
    model.eval()

    # 测试模型
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in dataloaders['test']:
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = correct / total
    print('Accuracy on test set: {:.4f}'.format(accuracy))

if __name__ == '__main__':
    model_path = r'D:\python\study\MobileNet V2\save'
    # 数据预处理和增强
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'test': transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    }

    # 加载数据集
    data_dir = r'D:\python\study\MobileNet V2\dataset_pz'
    image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
                      for x in ['train', 'test']}
    dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4, shuffle=True, num_workers=4)
                   for x in ['train', 'test']}
    class_names = image_datasets['train'].classes

    model = models.mobilenet_v2(pretrained=True)
    train(model_path, dataloaders, class_names)
    test(model_path, dataloaders, class_names)